#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
数据库配置模块
定义数据库连接和会话管理
"""

import os
from typing import Generator
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 数据库连接配置
# 从环境变量获取配置，如果不存在则使用默认值
# 使用用户提供的密码
DB_CONFIG = {
    "host": os.getenv("DB_HOST", "localhost"),
    "user": os.getenv("DB_USER", "root"),
    "password": os.getenv("DB_PASSWORD", "123456"),  # 用户提供的密码
    "database": os.getenv("DB_NAME", "trading"),
    "port": int(os.getenv("DB_PORT", "3306"))
}

# 构建数据库连接URL
DATABASE_URL = f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"

# 创建数据库引擎
# 参数说明：
# - echo: 是否打印SQL语句
# - pool_pre_ping: 连接池检查
# - pool_size: 连接池大小
# - max_overflow: 最大溢出连接数
engine = create_engine(
    DATABASE_URL,
    echo=False,  # 生产环境设置为False
    pool_pre_ping=True,
    pool_size=10,
    max_overflow=20,
    pool_recycle=3600,  # 连接回收时间（秒）
    pool_timeout=30  # 连接池超时时间（秒）
)

# 创建会话工厂
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 创建基类
Base = declarative_base()


def get_db() -> Generator[Session, None, None]:
    """
    获取数据库会话的依赖项
    用于FastAPI的依赖注入
    """
    db = SessionLocal()
    try:
        yield db
    except Exception as e:
        logger.error(f"数据库会话错误: {str(e)}")
        db.rollback()
        raise
    finally:
        db.close()


def init_db():
    """
    初始化数据库
    创建所有表
    """
    try:
        # 导入所有模型，确保它们被注册
        from src.models import user, product, transaction, wallet, message, review  # noqa
        
        # 创建所有表
        Base.metadata.create_all(bind=engine)
        logger.info("数据库表创建成功")
    except Exception as e:
        logger.error(f"数据库初始化失败: {str(e)}")
        raise


def check_db_connection():
    """
    检查数据库连接
    """
    try:
        # 测试连接
        with engine.connect() as conn:
            conn.execute("SELECT 1")
        logger.info("数据库连接成功")
        return True
    except Exception as e:
        logger.error(f"数据库连接失败: {str(e)}")
        return False


# 数据库事务装饰器
def transactional(func):
    """
    事务装饰器，用于自动处理事务
    """
    def wrapper(*args, **kwargs):
        db = SessionLocal()
        try:
            result = func(db, *args, **kwargs)
            db.commit()
            return result
        except Exception as e:
            db.rollback()
            logger.error(f"事务执行失败: {str(e)}")
            raise
        finally:
            db.close()
    return wrapper